import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import argparse

def load_dataset(data_dir):
    """加载预处理后的数据集"""
    samples = []
    labels = []
    for file in os.listdir(data_dir):
        if file.endswith('.npy'):
            data = np.load(os.path.join(data_dir, file), allow_pickle=True).item()
            samples.append(data['ecg'])
            labels.append(data['label'])
    return np.array(samples), np.array(labels)

def plot_single_sample(ecg_signal, label, sample_id, save_dir=None):
    """绘制单个ECG样本波形"""
    plt.figure(figsize=(12, 4))
    time = np.arange(len(ecg_signal)) / 512  # 假设采样率为512Hz
    plt.plot(time, ecg_signal, 'b-', linewidth=1)
    plt.title(f'ECG Signal (Label: {label})')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.grid(True)
    
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        plt.savefig(os.path.join(save_dir, f'sample_{sample_id}.png'))
        plt.close()
    else:
        plt.show()

def plot_label_distribution(labels, save_dir=None):
    """绘制标签分布统计图"""
    plt.figure(figsize=(10, 6))
    label_counts = Counter(labels)
    
    # 绘制柱状图
    sns.barplot(x=list(label_counts.keys()), y=list(label_counts.values()))
    plt.title('Label Distribution')
    plt.xlabel('Label')
    plt.ylabel('Count')
    
    # 添加数值标签
    for i, count in enumerate(label_counts.values()):
        plt.text(i, count, str(count), ha='center', va='bottom')
    
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        plt.savefig(os.path.join(save_dir, 'label_distribution.png'))
        plt.close()
    else:
        plt.show()

def main():
    parser = argparse.ArgumentParser(description='ECG数据可视化工具')
    parser.add_argument('--data_dir', type=str, required=True, help='预处理后的数据目录路径')
    parser.add_argument('--save_dir', type=str, default='results/plots', help='图像保存目录')
    parser.add_argument('--samples', type=int, default=5, help='要绘制的样本数量')
    parser.add_argument('--label', type=int, choices=[0, 1], default=None, 
                        help='指定要绘制的标签类型 (0 或 1)，不指定则绘制所有标签')
    args = parser.parse_args()

    # 加载数据
    print(f"正在从 {args.data_dir} 加载数据...")
    samples, labels = load_dataset(args.data_dir)
    print(f"加载完成，共 {len(samples)} 个样本")

    # 绘制标签分布
    plot_label_distribution(labels, args.save_dir)

    # 如果指定了标签，筛选对应的样本
    if args.label is not None:
        mask = labels == args.label
        samples = samples[mask]
        labels = labels[mask]
        print(f"筛选标签 {args.label} 的样本，共 {len(samples)} 个")

    # 绘制部分样本波形
    for i in range(min(args.samples, len(samples))):
        plot_single_sample(samples[i], labels[i], i, args.save_dir)

if __name__ == '__main__':
    main()